import asyncio
import socket
import json
import uuid
import time
import zmq.asyncio
from typing import Dict
from .common import *

class AsyncDeviceAgent:
    def __init__(self, name: str, info: str = "", pub_port: int = 6000):
        self.uuid = str(uuid.uuid4())
        self.name = name
        self.info = info
        self.pub_port = pub_port
        self.ip = get_local_ip()
        self.context = zmq.asyncio.Context()

    async def udp_broadcast(self):
        udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        udp_socket.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        message = json.dumps({
            "uuid": self.uuid,
            "name": self.name,
            "info": self.info,
            "pub_port": self.pub_port
        }).encode()

        while True:
            udp_socket.sendto(message, ('255.255.255.255', BROADCAST_PORT))
            await asyncio.sleep(BROADCAST_INTERVAL)

    async def zmq_heartbeat(self):
        pub_socket = self.context.socket(zmq.PUB)
        pub_socket.bind(f"tcp://*:{self.pub_port}")

        while True:
            msg = json.dumps({
                "uuid": self.uuid,
                "timestamp": time.time()
            })
            await pub_socket.send_string(msg)
            await asyncio.sleep(HEARTBEAT_INTERVAL)

    async def start(self):
        print(f"[Agent] 启动设备 {self.name}（UUID: {self.uuid}）")
        retry_count = 0
        while retry_count < MAX_RETRY_COUNT:
            try:
                await asyncio.gather(
                    self.udp_broadcast(),
                    self.zmq_heartbeat()
                )
            except Exception as e:
                retry_count += 1
                print(f"[错误] 设备运行异常: {e}")
                if retry_count < MAX_RETRY_COUNT:
                    print(f"[重试] {RECONNECT_DELAY}秒后进行第{retry_count}次重试...")
                    await asyncio.sleep(RECONNECT_DELAY)
                else:
                    print(f"[失败] 已达到最大重试次数({MAX_RETRY_COUNT})，设备停止运行")
                    break
            else:
                retry_count = 0  # 重置重试计数
                await asyncio.sleep(1)  # 避免过于频繁的重试

if __name__ == "__main__":
    agent = AsyncDeviceAgent(name="AlphaBot", info="Jetson Nano")
    asyncio.run(agent.start())
